import torch
import itertools
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks
from . import EnDecodes
from . import contraloss
import os


class HADIPREModel(BaseModel):
    """
    This class implements the CycleGAN model, for learning image-to-image translation without paired data.

    The model training requires '--dataset_mode unaligned' dataset.
    By default, it uses a '--netG resnet_9blocks' ResNet generator,
    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
    and a least-square GANs objective ('--gan_mode lsgan').

    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        parser.set_defaults(no_dropout=True)
        parser.add_argument('--feaLength', type=int, default=64, help='weight')
        return parser

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        self.loss_names = ['MSE']
        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        self.visual_names = ['Os','Bs','pred_Bs']

        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        if self.isTrain:
            self.model_names = ['E_Bs','G_B']
        else:  # during test time, only load Gs
            self.model_names = ['E_Bs']

        # define networks (both Generators and discriminators)
        # The naming is different from those used in the paper.
        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
        self.netE_Bs = EnDecodes.define_Encoder('encoder_content', opt.input_nc, opt.feaLength, self.gpu_ids)
        self.netG_B = EnDecodes.define_Encoder('decoder_single', opt.feaLength, 3, self.gpu_ids)

        if self.isTrain:  # define discriminators
            # define loss functions
            self.criterion_MSE = torch.nn.MSELoss()
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            self.optimizer_G = torch.optim.Adam(itertools.chain(self.netE_Bs.parameters(),
                                                                self.netG_B.parameters()),lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizers.append(self.optimizer_G)

    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        # AtoB = self.opt.direction == 'AtoB'
        # self.real_A = input['A' if AtoB else 'B'].to(self.device)
        # self.real_B = input['B' if AtoB else 'A'].to(self.device)
        # self.image_paths = input['A_paths' if AtoB else 'B_paths']

        self.Os = input['O_s'].to(self.device)
        self.Bs = input['B_s'].to(self.device)
        self.Or = input['O_t'].to(self.device)
        # if not self.isTrain:
        if self.opt.Bt_access:
            self.Bt = input['B_t'].to(self.device)
        # self.Bt = input['B_t'].to(self.device)
        self.image_paths = input['path']  # for test

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        # Image Space
        self.f_Bs = self.netE_Bs(self.Os)
        self.pred_Bs = self.netG_B(self.f_Bs)

    def backward_G(self):
        # Synthetic MSE loss
        self.loss_MSE = self.criterion_MSE(self.Bs, self.pred_Bs)
        # Total loss
        loss_G = self.loss_MSE
        loss_G.backward()


    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # # forward
        self.forward()      # compute fake images and reconstruction images.

        # G_A and G_B
        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero
        self.backward_G()             # calculate gradients for G_A and G_B
        self.optimizer_G.step()       # update G_A and G_B's weights

